{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/AIWintermuteAI/aXeleRate/blob/master/resources/aXeleRate_test_detector.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hS9yMrWe02WQ"
   },
   "source": [
    "## Detection model Training and Inference\n",
    "\n",
    "In this notebook we will use axelerate, Keras-based framework for AI on the edge, to quickly setup model training and then after training session is completed convert it to .tflite and .kmodel formats.\n",
    "\n",
    "First, let's take care of some administrative details. \n",
    "\n",
    "1) Before we do anything, make sure you have choosen GPU as Runtime type (in Runtime - > Change Runtime type).\n",
    "\n",
    "2) We need to mount Google Drive for saving our model checkpoints and final converted model(s). Press on Mount Google Drive button in Files tab on your left. \n",
    "\n",
    "In the next cell we clone axelerate Github repository and import it. \n",
    "\n",
    "**It is possible to use pip install or python setup.py install, but in that case you will need to restart the enironment.** Since I'm trying to make the process as streamlined as possibile I'm using sys.path.append for import."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "y07yAbYbjV2s"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "M\taxelerate/networks/common_utils/convert.py\r\n",
      "M\taxelerate/networks/yolo/backend/utils/augment.py\r\n",
      "Already on 'dev'\r\n",
      "Your branch is up to date with 'origin/dev'.\r\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "#!git clone https://github.com/AIWintermuteAI/aXeleRate.git\n",
    "!cd aXeleRate && git checkout dev\n",
    "import sys\n",
    "sys.path.append('aXeleRate')\n",
    "from axelerate import setup_training,setup_inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "5TBRMPZ83dRL"
   },
   "source": [
    "At this step you typically need to get the dataset. You can use !wget command to download it from somewhere on the Internet or !cp to copy from My Drive as in this example\n",
    "```\n",
    "!cp -r /content/drive/'My Drive'/pascal_20_segmentation.zip .\n",
    "!unzip --qq pascal_20_segmentation.zip\n",
    "```\n",
    "For this notebook small test dataset is already in axelerate/sample_datasets folder, so no need to download anything.\n",
    "\n",
    "Let's visualize our detection model test dataset. There are two images in train folder with corresponding annotations in PASCAL-VOC format in annotations folder.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_tpsgkGj7d79"
   },
   "outputs": [],
   "source": [
    "from axelerate.networks.yolo.backend.utils.augment import visualize_dataset\n",
    "\n",
    "visualize_dataset(img_folder='aXeleRate/sample_datasets/detector/imgs', ann_folder='aXeleRate/sample_datasets/detector/anns', img_size=None, jitter=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "S1oqdtbr7VLB"
   },
   "source": [
    "Next step is defining a config dictionary. Most lines are self-explanatory.\n",
    "\n",
    "Type is model frontend - Classifier, Detector or Segnet\n",
    "\n",
    "Architecture is model backend (feature extractor) \n",
    "\n",
    "- Full Yolo\n",
    "- Tiny Yolo\n",
    "- MobileNet1_0\n",
    "- MobileNet7_5 \n",
    "- MobileNet5_0 \n",
    "- MobileNet2_5 \n",
    "- SqueezeNet\n",
    "- VGG16\n",
    "- ResNet50\n",
    "\n",
    "For more information on anchors, please read here\n",
    "https://github.com/pjreddie/darknet/issues/568\n",
    "\n",
    "Labels are labels present in your dataset.\n",
    "IMPORTANT: Please, list all the labels present in the dataset.\n",
    "\n",
    "object_scale determines how much to penalize wrong prediction of confidence of object predictors\n",
    "\n",
    "no_object_scale determines how much to penalize wrong prediction of confidence of non-object predictors\n",
    "\n",
    "coord_scale determines how much to penalize wrong position and size predictions (x, y, w, h)\n",
    "\n",
    "class_scale determines how much to penalize wrong class prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Jw4q6_MsegD2"
   },
   "outputs": [],
   "source": [
    "config = {\n",
    "        \"model\":{\n",
    "            \"type\":                 \"Detector\",\n",
    "            \"architecture\":         \"MobileNet7_5\",\n",
    "            \"input_size\":           224,\n",
    "            \"anchors\":              [0.57273, 0.677385, 1.87446, 2.06253, 3.33843, 5.47434, 7.88282, 3.52778, 9.77052, 9.16828],\n",
    "            \"labels\":               [\"aeroplane\",\"person\",\"diningtable\",\"bottle\",\"bird\",\"bus\",\"boat\",\"cow\",\"sheep\",\"train\"],\n",
    "            \"coord_scale\" : \t\t1.0,\n",
    "            \"class_scale\" : \t\t1.0,\n",
    "            \"object_scale\" : \t\t5.0,\n",
    "            \"no_object_scale\" : \t1.0\n",
    "        },\n",
    "        \"weights\" : {\n",
    "            \"full\":   \t\t\t\t\"\",\n",
    "            \"backend\":   \t\t    \"imagenet\"\n",
    "        },\n",
    "        \"train\" : {\n",
    "            \"actual_epoch\":         1,\n",
    "            \"train_image_folder\":   \"aXeleRate/sample_datasets/detector/imgs\",\n",
    "            \"train_annot_folder\":   \"aXeleRate/sample_datasets/detector/anns\",\n",
    "            \"train_times\":          4,\n",
    "            \"valid_image_folder\":   \"aXeleRate/sample_datasets/detector/imgs_validation\",\n",
    "            \"valid_annot_folder\":   \"aXeleRate/sample_datasets/detector/anns_validation\",\n",
    "            \"valid_times\":          4,\n",
    "            \"valid_metric\":         \"mAP\",\n",
    "            \"batch_size\":           4,\n",
    "            \"learning_rate\":        1e-4,\n",
    "            \"saved_folder\":   \t\t\"detector\",\n",
    "            \"first_trainable_layer\": \"\",\n",
    "            \"augumentation\":\t\t\t\tTrue,\n",
    "            \"is_only_detect\" : \t\tFalse\n",
    "        },\n",
    "        \"converter\" : {\n",
    "            \"type\":   \t\t\t\t[\"k210\"]\n",
    "        }\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "kobC_7gd5mEu"
   },
   "source": [
    "Let's check what GPU we have been assigned in this Colab session, if any."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rESho_T70BWq"
   },
   "outputs": [],
   "source": [
    "from tensorflow.python.client import device_lib\n",
    "device_lib.list_local_devices()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "cWyKjw-b5_yp"
   },
   "source": [
    "Finally we start the training by passing config dictionary we have defined earlier to setup_training function. The function will start the training with Checkpoint, Reduce Learning Rate on Plateau and Early Stopping callbacks. After the training has stopped, it will convert the best model into the format you have specified in config and save it to the project folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "deYD3cwukHsj"
   },
   "outputs": [],
   "source": [
    "model_path = setup_training(config_dict=config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ypTe3GZI619O"
   },
   "source": [
    "After training it is good to check the actual perfomance of your model by doing inference on your validation dataset and visualizing results. This is exactly what next block does. Obviously since our model has only trained on a few images the results are far from stellar, but if you have a good dataset, you'll have better results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Edge TPU Converter ready\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/miniconda3/envs/notebook/lib/python3.7/site-packages/tensorflow_core/lite/python/lite.py:854: UserWarning: Property target_ops is deprecated, please use target_spec.supported_ops instead.\n",
      "  \"target_spec.supported_ops instead.\" % name)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/github/aXeleRate-dev/resources\n",
      "edgetpu_compiler --out_dir /home/ubuntu/github/aXeleRate-dev/resources /home/ubuntu/github/aXeleRate-dev/resources/YOLO_best_mAP.tflite\n",
      "Edge TPU Compiler version 2.1.302470888\n",
      "\n",
      "\n",
      "\n",
      "Model compiled successfully in 508 ms.\n",
      "\n",
      "\n",
      "\n",
      "Input model: /home/ubuntu/github/aXeleRate-dev/resources/YOLO_best_mAP.tflite\n",
      "\n",
      "Input size: 1.97MiB\n",
      "\n",
      "Output model: /home/ubuntu/github/aXeleRate-dev/resources/YOLO_best_mAP_edgetpu.tflite\n",
      "\n",
      "Output size: 2.03MiB\n",
      "\n",
      "On-chip memory used for caching model parameters: 1.96MiB\n",
      "\n",
      "On-chip memory remaining for caching model parameters: 5.85MiB\n",
      "\n",
      "Off-chip memory used for streaming uncached model parameters: 0.00B\n",
      "\n",
      "Number of Edge TPU subgraphs: 1\n",
      "\n",
      "Total number of operations: 36\n",
      "\n",
      "Operation log: /home/ubuntu/github/aXeleRate-dev/resources/YOLO_best_mAP_edgetpu.log\n",
      "\n",
      "\n",
      "\n",
      "Model successfully compiled but not all operations are supported by the Edge TPU. A percentage of the model will instead run on the CPU, which is slower. If possible, consider updating your model to use only operations supported by the Edge TPU. For details, visit g.co/coral/model-reqs.\n",
      "\n",
      "Number of operations that will run on Edge TPU: 34\n",
      "\n",
      "Number of operations that will run on CPU: 2\n",
      "\n",
      "See the operation log file for individual operation details.\n",
      "\n",
      "0\n"
     ]
    }
   ],
   "source": [
    "from axelerate.networks.common_utils.convert import Converter\n",
    "converter = Converter('edgetpu', 'MobileNet7_5', 'aXeleRate/sample_datasets/detector/imgs_validation')\n",
    "converter.convert_model('YOLO_best_mAP.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from axelerate.networks.common_utils.convert import Converter\n",
    "converter = Converter('k210', 'MobileNet7_5', 'aXeleRate/sample_datasets/detector/imgs_validation')\n",
    "converter.convert_model('YOLO_best_mAP.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from axelerate.networks.common_utils.convert import Converter\n",
    "converter = Converter('openvino', 'MobileNet7_5', 'aXeleRate/sample_datasets/detector/imgs_validation')\n",
    "converter.convert_model('YOLO_best_mAP.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jE7pTYmZN7Pi"
   },
   "outputs": [],
   "source": [
    "from keras import backend as K \n",
    "K.clear_session()\n",
    "setup_inference(config, model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "5YuVe2VD11cd"
   },
   "source": [
    "Good luck and happy training! Have a look at these articles, that would allow you to get the most of Google Colab or connect to local runtime if there are no GPUs available;\n",
    "\n",
    "https://medium.com/@oribarel/getting-the-most-out-of-your-google-colab-2b0585f82403\n",
    "\n",
    "https://research.google.com/colaboratory/local-runtimes.html"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyNfsQT8mfok4Na7Ben6VVvQ",
   "collapsed_sections": [],
   "include_colab_link": true,
   "mount_file_id": "1rCJbj9BGoDxEt1ERSK3onxShVBv9LS7B",
   "name": "aXeleRate_test_detector.ipynb",
   "private_outputs": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
